package com.hoppinzq.service.aop;

import com.hoppinzq.service.aop.annotation.RateLimit;
import com.hoppinzq.service.exception.ResultReturnException;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.stereotype.Component;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import javax.servlet.http.HttpServletRequest;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Aspect
@Component
public class RateLimitAspect {

    private Map<String, Map<String, Long>> limitMap = new ConcurrentHashMap<>();

    @Around("@annotation(rateLimit)")
    public Object limit(ProceedingJoinPoint point, RateLimit rateLimit) throws Throwable {
        // 获取限流阈值和时间间隔
        int limit = rateLimit.limit();
        long timeout = rateLimit.timeout();

        // 获取 HttpServletRequest 对象
        ServletRequestAttributes attributes = (ServletRequestAttributes) RequestContextHolder.getRequestAttributes();
        HttpServletRequest request = attributes.getRequest();

        // 获取客户端 IP 地址
        String ip = request.getRemoteAddr();

        // 获取方法名
        MethodSignature signature = (MethodSignature) point.getSignature();
        String methodName = signature.getMethod().getName();

        // 判断该方法是否需要进行限流
        if (limit <= 0 || timeout <= 0) {
            return point.proceed();
        }

        // 获取当前时间戳
        long now = System.currentTimeMillis();

        // 获取该 IP 在该方法上的访问记录
        Map<String, Long> ipMap = limitMap.get(methodName);
        if (ipMap == null) {
            ipMap = new ConcurrentHashMap<>();
            limitMap.put(methodName, ipMap);
        }
        Long last = ipMap.get(ip);

        // 如果上一次访问时间在限流时间间隔内，则进行限流
        if (last != null && now - last < timeout) {
            throw new ResultReturnException("请求频繁,请稍后调用", 403);
        }

        // 如果该 IP 在该方法上的访问记录不存在或者已经过期，则更新访问时间
        ipMap.put(ip, now);

        // 如果该 IP 在该方法上的访问记录数超过了限流阈值，则进行限流
        if (ipMap.size() > limit) {
            throw new ResultReturnException("请求频繁,请稍后调用", 403);
        }

        // 执行目标方法
        return point.proceed();
    }
}
